﻿using System.Text;
using System.Text.Json;
using SharpToken;

namespace OpenAIProxy;

// 流式响应捕获器中间件
// public class StreamingCaptureMiddleware
// {
//     private readonly RequestDelegate _next;
//     private readonly ILogger<StreamingCaptureMiddleware> _logger;
//     
//     public StreamingCaptureMiddleware(RequestDelegate next, ILogger<StreamingCaptureMiddleware> logger)
//     {
//         _next = next;
//         _logger = logger;
//     }
//
//     public async Task InvokeAsync(HttpContext context, TokenAccumulator tokenAccumulator)
//     {
//         // 仅处理流式请求
//         if (!context.Request.Path.StartsWithSegments("/v1/chat/completions") || 
//             !context.Request.HasJsonContentType() ||
//             context.Request.Method != "POST")
//         {
//             await _next(context);
//             return;
//         }
//
//         // 创建捕获缓冲区
//         var originalBody = context.Response.Body;
//         using var captureBuffer = new MemoryStream();
//         context.Response.Body = new TeeStream(originalBody, captureBuffer);
//         tokenAccumulator.Reset();
//
//         try
//         {
//             await _next(context);
//         }
//         finally
//         {
//             context.Response.Body = originalBody;
//             
//             // 仅当流式传输成功完成时处理内容
//             if (context.Response.StatusCode == 200)
//             {
//                 await ProcessCapturedContent(captureBuffer, tokenAccumulator);
//             }
//         }
//     }
//
//     private async Task ProcessCapturedContent(MemoryStream captureStream, TokenAccumulator tokenAccumulator)
//     {
//         captureStream.Seek(0, SeekOrigin.Begin);
//         var reader = new StreamReader(captureStream);
//         var chunks = ParseSSEChunks(await reader.ReadToEndAsync());
//         
//         foreach (var chunk in chunks)
//         {
//             if (!string.IsNullOrWhiteSpace(chunk.Content))
//             {
//                 tokenAccumulator.AddContent(chunk.Content);
//             }
//             _logger.LogInformation($"Captured chunk: {chunk.ToString()}");
//         }
//
//         _logger.LogInformation($"Total tokens: {tokenAccumulator.TotalTokens}");
//     }
//
//     private List<(string RawChunk, string? Content)> ParseSSEChunks(string sseContent)
//     {
//         var chunks = new List<(string, string?)>();
//         var lines = sseContent.Split('\n');
//         string? currentChunk = null;
//
//         for (int i = 0; i < lines.Length; i++)
//         {
//             var line = lines[i].Trim();
//             if (string.IsNullOrWhiteSpace(line)) continue;
//
//             if (line.StartsWith("data: "))
//             {
//                 currentChunk = line.Substring(5).Trim();
//                 
//                 // 如果下一个是空行，则完成一个块
//                 if (i + 1 < lines.Length && string.IsNullOrWhiteSpace(lines[i + 1]))
//                 {
//                     chunks.Add((currentChunk, ExtractContent(currentChunk)));
//                     currentChunk = null;
//                     i++; // 跳过空行
//                 }
//             }
//             else if (currentChunk != null)
//             {
//                 // 合并多行块
//                 currentChunk += $"\n{line}";
//             }
//         }
//
//         // 处理最后一个块
//         if (currentChunk != null)
//         {
//             chunks.Add((currentChunk, ExtractContent(currentChunk)));
//         }
//
//         return chunks;
//     }
//
//     private string? ExtractContent(string dataChunk)
//     {
//         try
//         {
//             if (dataChunk == "[DONE]") return null;
//             
//             var json = JsonDocument.Parse(dataChunk);
//             if (json.RootElement.TryGetProperty("choices", out var choices) && 
//                 choices.GetArrayLength() > 0 &&
//                 choices[0].TryGetProperty("delta", out var delta) &&
//                 delta.TryGetProperty("content", out var content))
//             {
//                 return content.GetString();
//             }
//         }
//         catch (JsonException)
//         {
//             // 非JSON数据块
//         }
//         return null;
//     }
// }

public class TeeStream : Stream
{
    private readonly Stream _primary;
    private readonly Stream _secondary;
    
    public TeeStream(Stream primary, Stream secondary)
    {
        _primary = primary;
        _secondary = secondary;
    }
    // 只重写异步方法，禁用同步方法
    public override async ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
    {
        await _primary.WriteAsync(buffer, cancellationToken);
        await _secondary.WriteAsync(buffer, cancellationToken);
    }
    
    public override async Task FlushAsync(CancellationToken cancellationToken)
    {
        await _primary.FlushAsync(cancellationToken);
        await _secondary.FlushAsync(cancellationToken);
    }
    
    // 禁用所有同步操作
    public override void Flush() => throw new InvalidOperationException("Synchronous operations are disallowed");
    public override void Write(byte[] buffer, int offset, int count) => throw new InvalidOperationException("Synchronous operations are disallowed");
    
    public override bool CanRead => false;
    public override bool CanSeek => false;
    public override bool CanWrite => true;
    public override long Length => throw new NotSupportedException();
    public override long Position { get => throw new NotSupportedException(); set => throw new NotSupportedException(); }
    public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
    public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
    public override void SetLength(long value) => throw new NotSupportedException();
}
// 修复的流式捕获中间件
public class StreamingCaptureMiddleware
{
    private readonly RequestDelegate _next;
    private readonly ILogger<StreamingCaptureMiddleware> _logger;
    
    public StreamingCaptureMiddleware(
        RequestDelegate next, 
        ILogger<StreamingCaptureMiddleware> logger)
    {
        _next = next;
        _logger = logger;
    }
    public async Task InvokeAsync(
        HttpContext context,
        TokenAccumulator tokenAccumulator,
        IAnalysisLogger analysisLogger)
    {
        // 只处理特定路由的流式请求
        if (!IsStreamingRequest(context))
        {
            await _next(context);
            return;
        }
        tokenAccumulator.Reset();
        var captureId = Guid.NewGuid().ToString();
        
        // 创建捕获缓存
        var captureBuffer = new MemoryStream();
        var originalBody = context.Response.Body;
        
        // 使用完全异步的TeeStream
        context.Response.Body = new TeeStream(originalBody, captureBuffer);
        try
        {
            await _next(context);
        }
        finally
        {
            context.Response.Body = originalBody;
            captureBuffer.Seek(0, SeekOrigin.Begin);
            
            if (context.Response.StatusCode == 200)
            {
                _ = ProcessContentAsync(captureBuffer, tokenAccumulator, captureId, analysisLogger);
            }
            else
            {
                _logger.LogWarning($"Streaming request failed with status {context.Response.StatusCode}");
            }
        }
    }
    
    private bool IsStreamingRequest(HttpContext context) 
        => context.Request.Path.StartsWithSegments("/v1/chat/completions", StringComparison.OrdinalIgnoreCase) &&
           context.Request.Method == "POST" &&
           context.Request.Query.TryGetValue("stream", out var streamParam) && 
           streamParam == "true";
    private async Task ProcessContentAsync(
        MemoryStream captureBuffer,
        TokenAccumulator tokenAccumulator,
        string captureId,
        IAnalysisLogger analysisLogger)
    {
        try
        {
            using var reader = new StreamReader(captureBuffer, Encoding.UTF8, leaveOpen: false);
            var sseContent = await reader.ReadToEndAsync();
            
            var chunks = ParseSSEChunks(sseContent);
            
            foreach (var chunk in chunks)
            {
                if (chunk.Type == ChunkType.Content && !string.IsNullOrEmpty(chunk.Data))
                {
                    tokenAccumulator.AddContent(chunk.Data);
                }
            }
            
            await analysisLogger.LogAsync(new StreamAnalysis(
                CaptureId: captureId,
                Timestamp: DateTime.UtcNow,
                TotalTokens: tokenAccumulator.TotalTokens,
                Chunks: chunks
            ));
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error processing captured stream content");
        }
    }
    
    private List<StreamChunk> ParseSSEChunks(string sseContent)
    {
        var chunks = new List<StreamChunk>();
        if (string.IsNullOrWhiteSpace(sseContent)) return chunks;
        var lines = sseContent.Split('\n');
        var currentChunk = new StringBuilder();
        for (int i = 0; i < lines.Length; i++)
        {
            var line = lines[i].Trim();
            if (string.IsNullOrWhiteSpace(line)) continue;
            if (line.StartsWith("data: "))
            {
                var chunkData = line.Substring(5).Trim();
                
                // 识别完成信号
                if (chunkData == "[DONE]")
                {
                    chunks.Add(new StreamChunk(
                        ChunkType.Done,
                        RawData: "[DONE]",
                        Content: null,
                        Data:""
                    ));
                    continue;
                }
                
                // 尝试解析内容块
                try
                {
                    var json = JsonDocument.Parse(chunkData);
                    if (json.RootElement.TryGetProperty("choices", out var choices) && 
                        choices.GetArrayLength() > 0)
                    {
                        if (choices[0].TryGetProperty("delta", out var delta) &&
                            delta.TryGetProperty("content", out var content))
                        {
                            var contentValue = content.GetString();
                            chunks.Add(new StreamChunk(
                                ChunkType.Content,
                                RawData: chunkData,
                                Content: contentValue,
                                ""
                            ));
                            continue;
                        }
                    }
                }
                catch (JsonException)
                {
                    // 不符合内容格式的记录为原始格式
                }
                
                // 无法解析的内容记录为元数据
                chunks.Add(new StreamChunk(
                    ChunkType.Metadata,
                    RawData: chunkData,
                    Content: null,
                    ""
                ));
            }
        }
        
        return chunks;
    }
}
// DTOs for analysis
public enum ChunkType { Content, Metadata, Done }
public record StreamChunk(
    ChunkType Type,
    string RawData,
    string? Content,
    string Data
);
public record StreamAnalysis(
    string CaptureId,
    DateTime Timestamp,
    int TotalTokens,
    List<StreamChunk> Chunks
);
// 分析日志服务接口
public interface IAnalysisLogger
{
    Task LogAsync(StreamAnalysis analysis);
}
// 简单的控制台日志分析实现
public class ConsoleAnalysisLogger : IAnalysisLogger
{
    public Task LogAsync(StreamAnalysis analysis)
    {
        var totalSize = Encoding.UTF8.GetByteCount(string.Join("", analysis.Chunks.Select(c => c.RawData)));
        
        Console.WriteLine($"\n==== Stream Analysis [{analysis.CaptureId}] ====");
        Console.WriteLine($"Timestamp: {analysis.Timestamp:O}");
        Console.WriteLine($"Token Count: {analysis.TotalTokens}");
        Console.WriteLine($"Chunk Count: {analysis.Chunks.Count}");
        Console.WriteLine($"Content Chunks: {analysis.Chunks.Count(c => c.Type == ChunkType.Content)}");
        Console.WriteLine($"Data Size: {totalSize} bytes\n");
        
        int index = 1;
        foreach (var chunk in analysis.Chunks)
        {
            Console.WriteLine($"[{index++}] {chunk.Type}:");
            if (chunk.Type == ChunkType.Content)
            {
                Console.WriteLine($"  Content: {chunk.Content}");
            }
            Console.WriteLine($"  Raw: {chunk.RawData}\n");
        }
        Console.WriteLine("==== End Analysis ====");
        
        return Task.CompletedTask;
    }
}
// 报名服务注册
public static class StreamingCaptureExtensions
{
    public static IServiceCollection AddStreamingCapture(this IServiceCollection services)
    {
        services.AddSingleton<TokenAccumulator>();
        services.AddSingleton<IAnalysisLogger, ConsoleAnalysisLogger>();
        services.AddTransient<ITokenCounter, SharpTokenCounter>();
        return services;
    }
    
    public static IApplicationBuilder UseStreamingCapture(this IApplicationBuilder app)
    {
        app.UseMiddleware<StreamingCaptureMiddleware>();
        return app;
    }
}

// 令牌累加器（简化版）
public class TokenAccumulator : ITokenAccumulator
{
    //     private readonly ITokenCounter _tokenCounter;
//     private readonly StringBuilder _contentBuilder = new();
//     private int? _tokenCount;
//
//     public TokenAccumulator(ITokenCounter tokenCounter)
//     {
//         _tokenCounter = tokenCounter;
//     }

    private readonly StringBuilder _contentBuilder = new();
    private readonly ITokenCounter _tokenCounter;
    private int _totalTokens;
    private int? _tokenCount;
    public TokenAccumulator(ITokenCounter tokenCounter)
    {
        _tokenCounter = tokenCounter;
    }
    
    public void Reset() => _totalTokens = 0;
    
    public void AddContent(string? content)
    {
        if (!string.IsNullOrWhiteSpace(content))
        {
            _totalTokens += _tokenCounter.CalculateTokens(content);
        }
    }
    
    public int TotalTokens => _totalTokens;
         public void Append(string dataLine)
     {
         var jsonText = dataLine["data: ".Length..];
         var jsonData = JsonDocument.Parse(jsonText);
         
         var content = jsonData.RootElement
             .GetProperty("choices")[0]
             .GetProperty("delta")
             .TryGetProperty("content", out var contentProp) 
             ? contentProp.GetString() 
             : null;
         
         if (!string.IsNullOrEmpty(content))
         {
             _contentBuilder.Append(content);
             
             _tokenCount = (_tokenCount ?? 0) + _tokenCounter.CalculateTokens(content);
         }
     }

     public (string Content, int TokenCount) GetAccumulatedResponse()
     {
         var content = _contentBuilder.ToString();
         return (content, _tokenCount ?? _tokenCounter.CalculateTokens(content));
     }
}

// // 在Startup中注册中间件
// public void Configure(IApplicationBuilder app, IWebHostEnvironment env)
// {
//     // 其他中间件...
//     app.UseMiddleware<StreamingCaptureMiddleware>();
//     // 其他中间件...
// }

public class SharpTokenCounter : ITokenCounter
{
    private static readonly GptEncoding _encoder = GptEncoding.GetEncoding("cl100k_base");
    
    public int CalculateTokens(string content)
    {
        if (string.IsNullOrWhiteSpace(content)) return 0;
        
        try
        {
            return _encoder.Encode(content).Count;
        }
        catch
        {
            // 回退到基于字符的近似计数
            return (int)(content.Length * 0.25); // 每4字符≈1个token
        }
    }
}